import random
from typing import Union, List

from atk.case_generator.generator.generate_types import GENERATOR_REGISTRY
from atk.case_generator.generator.base_generator import CaseGenerator
from atk.configs.case_config import InputCaseConfig, CaseConfig
from functools import reduce


@GENERATOR_REGISTRY.register("generate_aclnn_rope_with_sin_cos_cache")
class aclnnRopeWithSinCosCacheGenerator(CaseGenerator):

    def after_case_config(self, case_config: CaseConfig) -> CaseConfig:

        headSize = random.randint(1,128)
        numQHeads = random.randint(1,32)
        numKHeads = random.randint(1,32)
        maxSeqLen = random.randint(1,2000)
        rotary_dim = random.randint(1,128)

        if rotary_dim > headSize:
            rotary_dim = headSize

        headSize = ((headSize + 32 - 1) // 32) * 32
        if case_config.inputs[1].dtype == 'fp32':
            rotary_dim = ((rotary_dim + 32 - 1) // 32) * 32
        else:
            rotary_dim = ((rotary_dim + 16 - 1) // 16) * 16

        case_config.inputs[0].range_values = [-1, maxSeqLen]

        case_config.inputs[1].shape[0] = case_config.inputs[0].shape[0]
        case_config.inputs[1].shape[1] = numQHeads * headSize

        case_config.inputs[2].shape[0] = case_config.inputs[0].shape[0]
        case_config.inputs[2].shape[1] = numKHeads * headSize
        case_config.inputs[2].dtype = case_config.inputs[1].dtype

        case_config.inputs[3].shape[0] = maxSeqLen
        case_config.inputs[3].shape[1] = headSize
        case_config.inputs[3].dtype = case_config.inputs[1].dtype

        case_config.inputs[5].range_values = [headSize]

        return case_config  